package com.kfpanda.spring.batch.repository.map;

import java.util.Collection;
import java.util.Iterator;
import java.util.Map;

import org.springframework.batch.core.JobExecution;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.repository.dao.MapExecutionContextDao;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.support.SerializationUtils;
import org.springframework.batch.support.transaction.TransactionAwareProxyFactory;

public class MapExecutionContextDaoSupport extends MapExecutionContextDao implements com.kfpanda.spring.batch.repository.map.MapExecutionContextDao {
	
	protected Map<Long, ExecutionContext> contextsByStepExecutionId = TransactionAwareProxyFactory
			.createAppendOnlyTransactionalMap();

	protected Map<Long, ExecutionContext> contextsByJobExecutionId = TransactionAwareProxyFactory
			.createAppendOnlyTransactionalMap();

	public void clear() {
		contextsByJobExecutionId.clear();
		contextsByStepExecutionId.clear();
	}

	private static ExecutionContext copy(ExecutionContext original) {
		return (ExecutionContext) SerializationUtils.deserialize(SerializationUtils.serialize(original));
	}

	public ExecutionContext getExecutionContext(StepExecution stepExecution) {
		return copy(contextsByStepExecutionId.get(stepExecution.getId()));
	}

	public void updateExecutionContext(StepExecution stepExecution) {
		ExecutionContext executionContext = stepExecution.getExecutionContext();
		if (executionContext != null) {
			contextsByStepExecutionId.put(stepExecution.getId(), copy(executionContext));
		}
	}

	public ExecutionContext getExecutionContext(JobExecution jobExecution) {
		return copy(contextsByJobExecutionId.get(jobExecution.getId()));
	}

	public void updateExecutionContext(JobExecution jobExecution) {
		ExecutionContext executionContext = jobExecution.getExecutionContext();
		if (executionContext != null) {
			contextsByJobExecutionId.put(jobExecution.getId(), copy(executionContext));
		}
	}

	public void saveExecutionContext(JobExecution jobExecution) {
		updateExecutionContext(jobExecution);
	}

	public void saveExecutionContext(StepExecution stepExecution) {
		updateExecutionContext(stepExecution);
	}

	@Override
	public synchronized void remove(JobExecution jobExecution) {
		this.contextsByJobExecutionId.remove(jobExecution.getId());
		Collection<StepExecution> steps = jobExecution.getStepExecutions();
		for(Iterator<StepExecution> itr = steps.iterator(); itr.hasNext(); ){
			StepExecution step = itr.next();
			this.contextsByStepExecutionId.remove(step.getId());
		}
	}

}
